// Copyright (c) 2013, Hannes Würfel <hannes.wuerfel@student.hpi.uni-potsdam.de>
// Computer Graphics Systems Group at the Hasso-Plattner-Institute, Germany
// All rights reserved.
//
// Redistribution and use in source and binary forms, with or without
// modification, are permitted provided that the following conditions are met:
//   * Redistributions of source code must retain the above copyright notice,
//     this list of conditions and the following disclaimer.
//   * Redistributions in binary form must reproduce the above copyright
//     notice, this list of conditions and the following disclaimer in the
//     documentation and/or other materials provided with the distribution.
//   * Neither the name of the Computer Graphics Systems Group at the
//     Hasso-Plattner-Institute (HPI), Germany nor the names of its
//     contributors may be used to endorse or promote products derived from
//     this software without specific prior written permission.
//
// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
// AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
// IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
// ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
// LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
// CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
// SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
// INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
// CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
// ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
// POSSIBILITY OF SUCH DAMAGE.

#include "TruncatedABufferTransparency.h"
#include "InputLayout.h"
#include "EffectRegistry.h"
#include "RenderState.h"

TruncatedABufferTransparency::TruncatedABufferTransparency(unsigned int averageTranslucentOverdraw)
    : m_headPointerBufSRV(0),
      m_nodesBufSRV(0),
      m_headPointerBufUAV(0),
      m_nodesBufUAV(0),
      m_offscreenRTV(0),
      m_offscreenSRV(0),
      m_nodesBuf(0),
      m_headPointerBuf(0),
      m_averageTranslucentOverdraw(averageTranslucentOverdraw)
{

}

TruncatedABufferTransparency::~TruncatedABufferTransparency()
{
    releaseCOM(m_offscreenRTV);
    releaseCOM(m_offscreenSRV);
    releaseCOM(m_headPointerBufSRV);
    releaseCOM(m_headPointerBufUAV);
    releaseCOM(m_nodesBufSRV);
    releaseCOM(m_nodesBufUAV);
    releaseCOM(m_nodesBuf);
    releaseCOM(m_headPointerBuf);
    releaseCOM(m_quadIB);
    releaseCOM(m_quadVB);
}

void TruncatedABufferTransparency::init(ID3D11Device* device, int msaaQuality)
{
    createScreenQuad(device);
    m_4xMSAAQuality = msaaQuality;
}

void TruncatedABufferTransparency::render(ID3D11DeviceContext* context, ID3D11RenderTargetView* backBufferRTV, ID3D11DepthStencilView* depthStencilView, ID3D11ShaderResourceView* backgroundTex)
{
    // set render target, depth buffer and UAVs
    ID3D11RenderTargetView* rtvs[1] = {backBufferRTV};
    ID3D11UnorderedAccessView* uavs[4] = {NULL, NULL, NULL, NULL};
    unsigned int uavCounters[4] = {0, 0, 0, 0};

    // set rendering to backbuffer
    context->OMSetRenderTargetsAndUnorderedAccessViews(1, rtvs, NULL, 1, 4, uavs, uavCounters);

    // #TODO really no depth buffer bound?!
    // no depthstencil test since no depth buffer is bound
    context->OMSetDepthStencilState(RenderState::disableDepthTestDSS, 0);

    // manual blending in shader
    context->OMSetBlendState(0, 0, 0xffffffff);

    // render quad
    context->IASetInputLayout(InputLayout::PosNormalTex);
    context->IASetPrimitiveTopology(D3D11_PRIMITIVE_TOPOLOGY_TRIANGLELIST);

    unsigned int stride = sizeof(Vertex::PosNormalTex);
    unsigned int offset = 0;
    context->IASetVertexBuffers(0, 1, &m_quadVB, &stride, &offset);
    context->IASetIndexBuffer(m_quadIB, DXGI_FORMAT_R32_UINT, 0);

    XMMATRIX identity = XMMatrixIdentity();
    EffectRegistry::truncatedABufferFX->setWorld(identity);
    EffectRegistry::truncatedABufferFX->setWorldInvTranspose(identity);
    EffectRegistry::truncatedABufferFX->setWorldViewProjection(identity);
    EffectRegistry::truncatedABufferFX->setInputHeadPointerBuffer(m_headPointerBufSRV);
    EffectRegistry::truncatedABufferFX->setInputNodesBuffer(m_nodesBufSRV);
    EffectRegistry::truncatedABufferFX->setBackgroundTexture(backgroundTex);

    ID3DX11EffectTechnique* tech = EffectRegistry::truncatedABufferFX->sortFragmentsTech();

    D3DX11_TECHNIQUE_DESC techDesc;
    tech->GetDesc(&techDesc);
    for(unsigned int i = 0; i < techDesc.Passes; ++i)
    {
        tech->GetPassByIndex(i)->Apply(0, context);
        context->DrawIndexed(6, 0, 0);
    }
}

void TruncatedABufferTransparency::prepareStoreFragmentsPass(ID3D11DeviceContext* context, ID3D11DepthStencilView* depthStencilView)
{
    float bgColor[] = {0.75f, 0.75f, 0.75f, 1.0f};
    float blendFactor[] = {0.0f, 0.0f, 0.0f, 0.0f};

    // this would be needed if there were opaque geometry
    /*ID3D11RenderTargetView* renderTargets[1] = {m_offscreenRTV};
    context->OMSetRenderTargets(1, renderTargets, depthStencilView);

    context->ClearRenderTargetView(m_offscreenRTV, bgColor);
    context->ClearDepthStencilView(depthStencilView, D3D11_CLEAR_DEPTH | D3D11_CLEAR_STENCIL, 1.0f, 0);
    */

    const unsigned int clearheadPointerBuffer[1] = {0xffffffff}; // -1
    context->ClearUnorderedAccessViewUint(m_headPointerBufUAV, clearheadPointerBuffer);

    // set render target, depth buffer and UAVs
    ID3D11RenderTargetView* rtvs[1] = {NULL};
    ID3D11UnorderedAccessView* uavs[4] = {m_headPointerBufUAV, m_nodesBufUAV, NULL, NULL};
    unsigned int uavCounters[4] = {0, 0, 0, 0};

    context->OMSetRenderTargetsAndUnorderedAccessViews(1, rtvs, depthStencilView, 1, 4, uavs, uavCounters);

    // set stencil buffer to increment for each fragment (sugar)
    context->OMSetDepthStencilState(RenderState::enableDepthTestNoDepthWritesStencilWriteIncrementDSS, 0);

    // disable color writes
    context->OMSetBlendState(RenderState::disableColorWritesBS, 0, 0xffffffff);

    EffectRegistry::truncatedABufferFX->setOutputHeadPointerBuffer(m_headPointerBufUAV);
    EffectRegistry::truncatedABufferFX->setOutputNodesBuffer(m_nodesBufUAV);
}

void TruncatedABufferTransparency::finishStoreFragmentsPass(ID3D11DeviceContext* context)
{
    float blendFactor[] = {0.0f, 0.0f, 0.0f, 0.0f};
    context->OMSetBlendState(0, blendFactor, 0xffffffff);

    context->OMSetDepthStencilState(0, 0);
}

void TruncatedABufferTransparency::createResources(ID3D11Device* device, int width, int height)
{
    releaseCOM(m_offscreenRTV); // not used yet TODO
    releaseCOM(m_offscreenSRV);
    releaseCOM(m_headPointerBufSRV);
    releaseCOM(m_headPointerBufUAV);
    releaseCOM(m_nodesBufSRV);
    releaseCOM(m_nodesBufUAV);
    releaseCOM(m_nodesBuf);
    releaseCOM(m_headPointerBuf);

    // create offscreen render target
    D3D11_TEXTURE2D_DESC texDesc;
    D3D11_RENDER_TARGET_VIEW_DESC;
    
    texDesc.Width = width;
    texDesc.Height = height;
    texDesc.MipLevels = 1;
    texDesc.ArraySize = 1;
    texDesc.Format = DXGI_FORMAT_R16G16B16A16_FLOAT; // DXGI_FORMAT_R8G8B8A8_UNORM_SRGB
    texDesc.SampleDesc.Count = 1;
    texDesc.SampleDesc.Quality = 0;
    texDesc.Usage = D3D11_USAGE_DEFAULT;
    texDesc.BindFlags = D3D11_BIND_RENDER_TARGET | D3D11_BIND_SHADER_RESOURCE;
    texDesc.CPUAccessFlags = 0;
    texDesc.MiscFlags = 0;

    ID3D11Texture2D* offscreenTex = 0;
    HANDLE_ERROR(device->CreateTexture2D(&texDesc, 0, &offscreenTex));
    HANDLE_ERROR(device->CreateShaderResourceView(offscreenTex, 0, &m_offscreenSRV));
    HANDLE_ERROR(device->CreateRenderTargetView(offscreenTex, 0, &m_offscreenRTV));

    releaseCOM(offscreenTex);

    // create structured node buffer
    D3D11_BUFFER_DESC bufDesc;
    D3D11_SHADER_RESOURCE_VIEW_DESC srvDesc;
    D3D11_UNORDERED_ACCESS_VIEW_DESC uavDesc;

    bufDesc.CPUAccessFlags = 0;
    bufDesc.BindFlags = D3D11_BIND_SHADER_RESOURCE | D3D11_BIND_UNORDERED_ACCESS;
    bufDesc.ByteWidth = static_cast<unsigned int>(m_averageTranslucentOverdraw * width * height * sizeof(NodeData));
    bufDesc.MiscFlags = D3D11_RESOURCE_MISC_BUFFER_STRUCTURED;
    bufDesc.Usage = D3D11_USAGE_DEFAULT;
    bufDesc.StructureByteStride = sizeof(NodeData);
    HANDLE_ERROR(device->CreateBuffer(&bufDesc, 0, &m_nodesBuf));

    uavDesc.Format = DXGI_FORMAT_UNKNOWN;
    uavDesc.ViewDimension = D3D11_UAV_DIMENSION_BUFFER;
    uavDesc.Buffer.FirstElement = 0;
    uavDesc.Buffer.NumElements = static_cast<unsigned int>(m_averageTranslucentOverdraw * width * height);
    uavDesc.Buffer.Flags = D3D11_BUFFER_UAV_FLAG_COUNTER; // needed for incrementCounter in shader
    HANDLE_ERROR(device->CreateUnorderedAccessView(m_nodesBuf, &uavDesc, &m_nodesBufUAV));

    srvDesc.Format = DXGI_FORMAT_UNKNOWN;
    srvDesc.ViewDimension = D3D11_SRV_DIMENSION_BUFFER;
    srvDesc.Buffer.ElementOffset = 0;
    srvDesc.Buffer.ElementWidth = static_cast<unsigned int>(m_averageTranslucentOverdraw * width * height);
    HANDLE_ERROR(device->CreateShaderResourceView(m_nodesBuf, &srvDesc, &m_nodesBufSRV));

    // create byte address head pointer buffer
    D3D11_BUFFER_DESC headPointerBufDesc;
    D3D11_SHADER_RESOURCE_VIEW_DESC srvHeadPointerDesc;
    D3D11_UNORDERED_ACCESS_VIEW_DESC uavHeadPointerDesc;

    headPointerBufDesc.BindFlags = D3D11_BIND_SHADER_RESOURCE | D3D11_BIND_UNORDERED_ACCESS;
    headPointerBufDesc.ByteWidth = sizeof(unsigned int) * width * height;
    headPointerBufDesc.MiscFlags = D3D11_RESOURCE_MISC_BUFFER_ALLOW_RAW_VIEWS;
    headPointerBufDesc.Usage = D3D11_USAGE_DEFAULT;
    headPointerBufDesc.CPUAccessFlags = 0;
    headPointerBufDesc.StructureByteStride = 0;
    HANDLE_ERROR(device->CreateBuffer(&headPointerBufDesc, 0, &m_headPointerBuf));

    uavHeadPointerDesc.Format = DXGI_FORMAT_R32_TYPELESS;
    uavHeadPointerDesc.ViewDimension = D3D11_UAV_DIMENSION_BUFFER;
    uavHeadPointerDesc.Buffer.FirstElement = 0;
    uavHeadPointerDesc.Buffer.NumElements = static_cast<unsigned int>(width * height);
    uavHeadPointerDesc.Buffer.Flags = D3D11_BUFFER_UAV_FLAG_RAW;
    HANDLE_ERROR(device->CreateUnorderedAccessView(m_headPointerBuf, &uavHeadPointerDesc, &m_headPointerBufUAV));

    srvHeadPointerDesc.Format = DXGI_FORMAT_R32_UINT;
    srvHeadPointerDesc.ViewDimension = D3D11_SRV_DIMENSION_BUFFER;
    srvHeadPointerDesc.Buffer.ElementOffset = 0;
    srvHeadPointerDesc.Buffer.ElementWidth = static_cast<unsigned int>(width * height);
    HANDLE_ERROR(device->CreateShaderResourceView(m_headPointerBuf, &srvHeadPointerDesc, &m_headPointerBufSRV));
}

void TruncatedABufferTransparency::createScreenQuad(ID3D11Device* device)
{
    GeometryProducer producer;
    //producer.createTestQuad(m_screenQuad);
    producer.createScreenQuad(m_screenQuad);

    D3D11_BUFFER_DESC vbd;
    vbd.Usage = D3D11_USAGE_IMMUTABLE;
    vbd.ByteWidth = sizeof(Vertex::PosNormalTex) * static_cast<unsigned int>(m_screenQuad.vertices.size());
    vbd.BindFlags = D3D11_BIND_VERTEX_BUFFER;
    vbd.CPUAccessFlags = 0;
    vbd.MiscFlags = 0;
    D3D11_SUBRESOURCE_DATA vinitData;
    vinitData.pSysMem = &m_screenQuad.vertices[0];
    HANDLE_ERROR(device->CreateBuffer(&vbd, &vinitData, &m_quadVB));

    D3D11_BUFFER_DESC ibd;
    ibd.Usage = D3D11_USAGE_IMMUTABLE;
    ibd.ByteWidth = sizeof(unsigned int) * static_cast<unsigned int>(m_screenQuad.indices.size());
    ibd.BindFlags = D3D11_BIND_INDEX_BUFFER;
    ibd.CPUAccessFlags = 0;
    ibd.MiscFlags = 0;
    D3D11_SUBRESOURCE_DATA iinitData;
    iinitData.pSysMem = &m_screenQuad.indices[0];
    HANDLE_ERROR(device->CreateBuffer(&ibd, &iinitData, &m_quadIB));
}